import torch
from torch import nn

from decoder import MaxViTDecoder
from encoder import MaxViT


class MaxVitUnet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1):
        super(MaxVitUnet, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.encoder_model = MaxViT(
            in_channels=self.in_channels,
            depths=(2, 2, 5, 2),
            channels=(64, 128, 256, 512),
            embed_dim=64,
            grid_window_size=(8, 8),
            attn_drop=0.2,
            drop=0.2,
            drop_path=0.2,
            debug=False,
        )
        self.decoder_model = MaxViTDecoder(
            in_channels=(64, 128, 256, 512),
            depths=(2, 2, 2),
            grid_window_size=(8, 8),
            attn_drop=0.2,
            drop=0.2,
            drop_path=0.2,
            num_classes=self.out_channels,  # 根据您的任务设置类别数量
            debug=False,
        )

    def forward(self, x):
        x = self.encoder_model(x)
        x = self.decoder_model(x)
        return nn.Sigmoid()(x)


if __name__ == "__main__":
    input = torch.randn(10, 3, 256, 256)
    model = MaxVitUnet(in_channels=3, out_channels=1)
    output = model(input)
    print(output.shape)
    print(f'Total parameters: {sum(param.numel() for param in model.parameters()) / 1e6:.2f} M')
